package x.ovo.wechat.bot;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.ReUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.time.Instant;
import java.util.regex.Pattern;

@Slf4j(topic = "AI绘图")
public class DrawService {

    private static final String GENERATE_URL = "https://image.baidu.com/aigc/generate";
    private static final String QUERY_URL = "https://image.baidu.com/aigc/query";
    private static final String DEFAULT_WIDTH = "512";
    private static final String DEFAULT_HEIGHT = "512";
    private static final Pattern PATTERN = Pattern.compile("(\\S+)\\s?(\\d+)?[\\s|x]?(\\d+)?");

    public static File draw(File floder, String content) {
        String prompt = ReUtil.get(PATTERN, content, 1);
        Integer height = Integer.parseInt(StrUtil.blankToDefault(ReUtil.get(PATTERN, content, 2), DEFAULT_HEIGHT));
        Integer width = Integer.parseInt(StrUtil.blankToDefault(ReUtil.get(PATTERN, content, 3), DEFAULT_WIDTH));
        JSONObject object = JSONUtil.createObj()
                .set("querycate", 10)
                .set("query", prompt)
                .set("width", width)
                .set("height", height);
        String taskid = JSONUtil.parseObj(HttpUtil.post(GENERATE_URL, object)).getStr("taskid");
        // 轮询，每隔500ms查询一次生成结果
        for (int i = 0; i < 20; i++) {
            ThreadUtil.safeSleep(500);
            JSONObject result = JSONUtil.parseObj(HttpUtil.get(QUERY_URL + "?taskid=" + taskid));
            if (result.getBool("isGenerate") && result.containsKey("picArr")) {
                String src = result.getByPath("picArr[0].src", String.class);
                File file = FileUtil.file(floder, Instant.now().toEpochMilli() + ".jpg");
                HttpUtil.downloadFile(src, file);
                return file;
            }
        }
        throw new RuntimeException("生成失败");
    }
}
